#!/opt/local/bin/python

from scipy import nan, isnan, pi, exp, sqrt, arctan
import matplotlib.pyplot as plt
import csv
from rfvectors import cavity, scale_roverq
from diffeqns import transientcavity, analysis, calcCurrents
rad2deg = lambda x: x*180/pi

class specificcavity(object):
    def __init__(self, filename, cavnum, I_b0, I_b=nan, delF=nan, Eampfile=nan, Ql=nan):
        self.I_b0 = I_b0
        if isnan(I_b):
            self.I_b = self.I_b0
        else:
            self.I_b = I_b
        self.delF = delF
        self.Eampfile = Eampfile

        data = csv.reader(open(filename, 'rU'), dialect=csv.excel_tab, delimiter=',')
        counter = 0
        f,V,beta = [],[],[]
        beamph,RoverQ,beta_0 = [],[],[]
        for line in data:
            f.append(float( line[4] ))
            V.append(float( line[0] ) * float( line[3] ))
            beta.append(float( line[2] ))
            beamph.append(float( line[1] ))
            RoverQ.append(float( line[5] ))
            beta_0.append(float( line[6] ))
        self.f = f[cavnum]
        self.V = V[cavnum]
        self.beta = beta[cavnum]
        self.beamph = beamph[cavnum]
        self.RoverQ = RoverQ[cavnum]
        self.beta_0 = beta_0[cavnum]
    
        if isnan(Ql):
            perfcav = cavity(f=self.f, delF=self.delF, I_b0=self.I_b0, I_b=self.I_b,
                    V_cav=self.V, beamph=self.beamph, RoverQ=self.RoverQ, beta_0=self.beta_0,
                    beta=self.beta, Eampfile=self.Eampfile)
            self.Ql = perfcav.Ql

    def steadystate(self):
        self.sscav = cavity(f=self.f, delF=self.delF, I_b0=self.I_b0, I_b=self.I_b, V_cav=self.V, beamph=self.beamph,
                RoverQ=self.RoverQ, Ql=self.Ql, beta_0=self.beta_0, beta=self.beta, Eampfile=self.Eampfile)
        self.RoverQscaling = scale_roverq(beta_0=self.sscav.beta_0, file=self.sscav.Eampfile)(beta=self.sscav.beta, freq=self.sscav.f)
        try:
            self.RoverQscaling = self.RoverQscaling[0]
        except IndexError:
            pass

    def drawss(self):
        self.ssfig    = plt.figure(1)
        self.sscav.drawvecs(self.ssfig)

    def transient(self, timestop, pulselen, acc):
        self.pulselen = pulselen
        self.steadystate()
        cav = transientcavity(freq=self.f, QL=self.sscav.Ql, detuning=self.sscav.delF)
        self.cavanal = analysis(tcavobj=cav, timestop=timestop, acc=acc)

        self.Ig,self.Ibeam,self.tinj = calcCurrents(cavobj=self.sscav, tcavobj=cav,
                time=self.cavanal.time, pulselength=pulselen, RoverQscaling=self.RoverQscaling)
        self.I = self.Ig + self.Ibeam

        self.cavanal.run(iniV=0, I=self.I)
        self.Vinc  = 0.5 * self.Ig * exp(cav.phi*1j) * cav.RL
        self.Pinc  = abs(self.Vinc)**2 / cav.RL
        self.Vref  = abs(self.Vinc) - sqrt(self.cavanal.V[:,0]**2 + self.cavanal.V[:,1]**2)
        self.Pref  = self.Vref**2 / cav.RL
        self.Vbeam = 0.5 * self.Ibeam * exp(cav.phi*1j) * cav.RL
        self.Pbeam = abs(self.Vbeam)**2 / cav.RL


    def drawtrans(self):
        self.transfig1 = plt.figure(11)
        self.transfig2 = plt.figure(12)
        self.transfig3 = plt.figure(13)
        self.transfig4 = plt.figure(14)

        ax1fig2 = self.transfig2.add_subplot(211)
        ax1fig2.plot(1e3*self.cavanal.time, 1e-6*self.Vinc.real, 'b', label='Real voltage')
        ax1fig2.plot(1e3*self.cavanal.time, 1e-6*self.Vinc.imag, 'r', label='Imag voltage')
        ax1fig2.plot(1e3*self.cavanal.time, 1e-6*abs(self.Vbeam), 'k', label='Abs beam voltage')
        ax1fig2.plot(1e3*self.cavanal.time, 1e-6*abs(self.Vinc), 'g', label='Abs voltage')
        ax1fig2.set_ylabel('Voltage / MV')
        
        ax2fig2 = self.transfig2.add_subplot(212)
        ax2fig2.plot(1e3*self.cavanal.time, abs(self.Vref)/abs(self.Vinc), label='Reflection coeff.')
        ax2fig2.set_ylabel('Reflection coefficient')
        
        ax1 = self.transfig1.add_subplot(211)
        ax2 = self.transfig1.add_subplot(212)
        ax1.plot(1e3*self.cavanal.time, 1e-6*self.cavanal.V[:,0], '-', label='Real cavity voltage')
        ax1.plot(1e3*self.cavanal.time, 1e-6*self.cavanal.V[:,1], '-', label='Imag cavity voltage')
        ax2.plot(1e3*self.cavanal.time, 1e-6*sqrt(self.cavanal.V[:,0]**2 + self.cavanal.V[:,1]**2), '-', label='Absolute cavity voltage')
        ax2.plot(1e3*self.cavanal.time, 1e-6*self.Vref, 'r', label='Reflected voltage')
        
        ax1fig3 = self.transfig3.add_subplot(111)
        ax1fig3.plot(1e3*self.cavanal.time, rad2deg(arctan(self.cavanal.V[:,1]/self.cavanal.V[:,0])), label='RF phase')
        ax1fig3.set_ylabel('RF phase (degrees)')
        
        ax1fig4 = self.transfig4.add_subplot(111)
        ax1fig4.plot(1e3*self.cavanal.time, self.Pref/1e3, label='Reflected: %0.3f J' % (sum(self.Pref * self.cavanal.timestep)))
        ax1fig4.set_ylabel('Power (kW)')
        
        for a in [ax1,ax2,ax1fig2,ax2fig2,ax1fig3,ax1fig4]:
            a.grid()
            a.set_xlabel('Time / ms')
            a.set_xlim(right=1e3*self.cavanal.time[-1])
            a.axvline(x=self.tinj*1e3, color='g')
            a.axvline(x=(self.tinj+self.pulselen)*1e3, color='g')
            a.legend(loc=0)
        ax1fig2.legend(loc=1)
        ax1.set_ylabel('Voltage / MV')
        ax2.set_ylabel('Abs voltage / MV')

if __name__=="__main__":
    A = specificcavity(filename='highbetanew.csv', cavnum=-1, I_b0=50e-3)
    #A.steadystate()
    #A.drawss()

    A.transient(timestop=5e-3, pulselen=2.8e-3, acc=1e-3)
    A.drawtrans()

    plt.show()

